from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import MessagesState
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.memory import MemorySaver
from langchain_core.tools import tool
from langgraph.prebuilt import ToolNode, create_react_agent
from utils import websearch, retriever
from typing import Annotated
import uuid


@tool
def chatbot_search(query: Annotated[str, "需要查询的问题描述"]):
    """可以使用chatbot来搜索互联网信息"""
    return websearch(query)


@tool
def chatbot_retrieve(query: Annotated[str, "需要查询的人物的信息"]):
    """如果用户查询的信息是关于人物的信息，使用chatbot_retrieve工具"""
    return retriever(query)


class ChatBot:

    def __init__(self):

        self._tools = [chatbot_search, chatbot_retrieve]

        self._llm = ChatOpenAI(
            api_key="ollama",
            model="qwen2.5:7b",
            base_url="http://192.168.10.13:60001/v1",
            temperature=0.6,
        ).bind_tools(self._tools)

        self._graph = self._init_graph()

    def _init_graph(self):
        _builder = StateGraph(MessagesState)
        _builder.add_node("chat_node", self._chat_agent)
        _builder.add_node("tool_node", ToolNode(tools=self._tools))

        _builder.add_edge(START, "chat_node")
        _builder.add_conditional_edges("chat_node", self._choose_tool_or_end_edge)
        _builder.add_edge("tool_node", "chat_node")

        _memory = MemorySaver()
        return _builder.compile(checkpointer=_memory)

    def _chat_agent(self, state):
        _message = self._llm.invoke(state["messages"])
        return {"messages": [_message]}

    def _choose_tool_or_end_edge(self, state):
        _last_message = state["messages"][-1]

        if _last_message.tool_calls:
            return "tool_node"
        else:
            return END

    def __call__(self):

        _thread_id = uuid.uuid4()
        while True:
            _human = input("我：")
            if _human == "bye":
                print("AI:再见！")
                break
            _config = {"thread_id": _thread_id}
            _ai = self._graph.invoke(
                {"messages": [HumanMessage(_human)]}, config=_config
            )
            print(f"AI：{_ai['messages'][-1].content}")


if __name__ == "__main__":
    chatbot = ChatBot()
    chatbot()
